package com.os.opencv.java.chapter12;

import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.TermCriteria;
import org.opencv.ml.RTrees;
import org.opencv.ml.Ml;
import org.opencv.ml.TrainData;

public class Rtrees {

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);

        //训练数据，每组两个数据，分别表示身高和体重
        float[] trainData = {180,81,178,75,185,87,168,65,
                173,68,160,52,158,55,170,63,163,60,168,66};
        Mat trainMat = new Mat(10,2, CvType.CV_32FC1);
        trainMat.put(0,0,trainData);

        //训练标签数据，0为男性，1为女性
        float[] label = {0,0,0,0,0,1,1,1,1,1};
        Mat labelMat = new Mat(10,1,CvType.CV_32FC1);
        labelMat.put(0,0,label);

        //测试数据
        float[] testData = {185,79,159,52};
        Mat testMat = new Mat(2,2,CvType.CV_32FC1);
        testMat.put(0,0,testData);

        //创建dtrees对象并设置参数
        RTrees rtrees = RTrees.create();
        rtrees.setMaxDepth(4);
        rtrees.setMinSampleCount(2);
        rtrees.setRegressionAccuracy(0.f);
        rtrees.setUseSurrogates(false);
        rtrees.setMaxCategories(16);
        rtrees.setPriors(new Mat());
        rtrees.setCalculateVarImportance(false);
        rtrees.setActiveVarCount(1);
        rtrees.setTermCriteria(new TermCriteria(TermCriteria.MAX_ITER, 5, 0));

        //训练随机森林模型并输出训练结果
        TrainData td = TrainData.create(trainMat, Ml.ROW_SAMPLE, labelMat);
        boolean result = rtrees.train(td.getSamples(), Ml.ROW_SAMPLE, td.getResponses());

        if(result){
            System.out.println("training successed!");
        }else{
            System.out.println("training failed!");
        }
        System.out.println();

        //保存决策树模型
        rtrees.save("rtree.xml");

        //对测试数据进行预测并输出预测结果
        Mat predicts = new Mat();
        rtrees.predict(testMat, predicts, 0);
        System.out.println("predicts:\n" + predicts.dump());
        System.out.println();

        //用文字输出判断结果
        for(int i=0; i<predicts.height(); i++){
            int predict = (int)predicts.get(i, 0)[0];
            if(predict == 0){
                System.out.println("No." + i + " is male");
            }else if(predict == 1){
                System.out.println("No." + i + " is female");
            }
        }

    }
    
}
